"""
股票数据采集器模块

负责通过akshare API获取股票数据，进行清洗，并保存到数据库
"""
import time
import datetime
import logging
import akshare as ak
import pandas as pd
from sqlalchemy.orm import Session
from sqlalchemy.exc import IntegrityError

from modules.models.base import get_db
from modules.models.stock import SecurityInfo, DailyQuote, CollectionTask
from config.settings import DATA_COLLECTION

# 配置日志
logger = logging.getLogger("stock_collector")


class StockCollector:
    """股票数据采集类"""
    
    def __init__(self):
        """初始化采集器"""
        self.retry_limit = DATA_COLLECTION["retry_limit"]
        self.retry_delay = DATA_COLLECTION["retry_delay"]
        self.batch_size = DATA_COLLECTION["batch_size"]
        self.db = next(get_db())
    
    def __del__(self):
        """析构函数，确保关闭数据库连接"""
        if hasattr(self, 'db') and self.db:
            self.db.close()
    
    def get_all_stock_list(self):
        """获取所有A股股票列表"""
        try:
            logger.info("开始获取A股股票列表...")
            # 使用akshare获取A股股票基本信息
            stock_info_df = ak.stock_info_a_code_name()
            logger.info(f"成功获取{len(stock_info_df)}只A股股票信息")
            
            # 添加港股和美股支持
            try:
                hk_stock_df = ak.stock_info_hk_code()
                logger.info(f"成功获取{len(hk_stock_df)}只港股信息")
                stock_info_df = pd.concat([stock_info_df, hk_stock_df])
            except Exception as e:
                logger.warning(f"获取港股列表失败: {str(e)}")
                
            return stock_info_df
        except Exception as e:
            logger.error(f"获取股票列表失败: {str(e)}")
            raise
    
    def collect_stock_info(self):
        """采集并更新股票基本信息"""
        logger.info("开始采集股票基本信息...")
        try:
            # 获取股票列表
            stock_df = self.get_all_stock_list()
            
            # 获取现有股票代码
            existing_codes = {
                code[0] for code in 
                self.db.query(SecurityInfo.code).filter(SecurityInfo.type == 'stock').all()
            }
            
            # 准备批量插入或更新的数据
            new_stocks = []
            update_stocks = []
            
            for _, row in stock_df.iterrows():
                code = row['code']
                name = row['name']
                
                # 根据代码前缀判断交易所
                if code.startswith('6'):
                    exchange = 'SSE'  # 上海证券交易所
                elif code.startswith(('0', '3')):
                    exchange = 'SZSE'  # 深圳证券交易所
                else:
                    exchange = 'OTHER'
                
                # 检查是否存在
                if code in existing_codes:
                    # 更新现有股票
                    stock = self.db.query(SecurityInfo).filter(
                        SecurityInfo.code == code,
                        SecurityInfo.type == 'stock'
                    ).first()
                    
                    if stock:
                        stock.name = name
                        stock.update_time = datetime.datetime.now()
                        update_stocks.append(stock)
                else:
                    # 创建新股票记录
                    new_stock = SecurityInfo(
                        code=code,
                        name=name,
                        type='stock',
                        exchange=exchange,
                        is_active=1
                    )
                    new_stocks.append(new_stock)
            
            # 批量保存新股票
            if new_stocks:
                self.db.add_all(new_stocks)
                self.db.commit()
                logger.info(f"成功添加{len(new_stocks)}只新股票")
            
            # 批量更新现有股票
            if update_stocks:
                self.db.commit()
                logger.info(f"成功更新{len(update_stocks)}只股票信息")
            
            return len(new_stocks), len(update_stocks)
        
        except Exception as e:
            self.db.rollback()
            logger.error(f"采集股票基本信息失败: {str(e)}")
            raise
    
    def collect_daily_quote(self, security_code, start_date=None, end_date=None, task_id=None):
        """
        采集指定股票的日线数据
        
        Args:
            security_code: 股票代码
            start_date: 开始日期，默认为None，表示从上次采集的位置继续
            end_date: 结束日期，默认为None，表示采集到当前日期
            task_id: 任务ID，用于更新任务状态
        
        Returns:
            新增数据条数
        """
        # 更新任务状态为运行中
        if task_id:
            self._update_task_status(task_id, "running")
        
        # 获取证券信息
        security = self.db.query(SecurityInfo).filter(
            SecurityInfo.code == security_code,
            SecurityInfo.type == 'stock'
        ).first()
        
        if not security:
            error_msg = f"股票代码 {security_code} 不存在"
            logger.error(error_msg)
            if task_id:
                self._update_task_status(task_id, "failed", error_msg)
            return 0
        
        # 如果未指定开始日期，则查找最新数据日期作为开始日期
        if not start_date:
            last_record = self.db.query(DailyQuote).filter(
                DailyQuote.security_id == security.id
            ).order_by(DailyQuote.trade_date.desc()).first()
            
            if last_record:
                # 从最后一条记录的下一天开始
                start_date = (last_record.trade_date + datetime.timedelta(days=1)).strftime('%Y%m%d')
            else:
                # 默认采集近5年数据
                start_date = (datetime.datetime.now() - datetime.timedelta(days=365*5)).strftime('%Y%m%d')
        
        # 如果未指定结束日期，则使用当前日期
        if not end_date:
            end_date = datetime.datetime.now().strftime('%Y%m%d')
        
        logger.info(f"开始采集股票 {security_code} ({security.name}) 的日线数据，时间范围: {start_date} - {end_date}")
        
        try:
            # 使用akshare获取日线数据
            retry_count = 0
            while retry_count < self.retry_limit:
                try:
                    # 处理不同交易所的股票代码格式
                    if security.exchange == 'SSE':
                        ak_code = f"{security_code}.SH"
                    elif security.exchange == 'SZSE':
                        ak_code = f"{security_code}.SZ"
                    else:
                        ak_code = security_code
                    
                    # 调用akshare API
                    df = ak.stock_zh_a_hist(symbol=security_code, period="daily", 
                                           start_date=start_date, end_date=end_date, 
                                           adjust="qfq")
                    break
                except Exception as e:
                    retry_count += 1
                    logger.warning(f"采集股票 {security_code} 数据失败，重试 {retry_count}/{self.retry_limit}: {str(e)}")
                    if retry_count >= self.retry_limit:
                        raise
                    time.sleep(self.retry_delay)
            
            if df.empty:
                logger.info(f"股票 {security_code} 在指定时间范围内没有数据")
                if task_id:
                    self._update_task_status(task_id, "success", None, 0)
                return 0
            
            # 转换和清洗数据
            df = self._clean_stock_data(df)
            
            # 保存到数据库
            inserted_count = self._save_daily_quotes(security.id, df)
            
            # 更新任务状态
            if task_id:
                self._update_task_status(task_id, "success", None, inserted_count)
            
            return inserted_count
            
        except Exception as e:
            error_msg = f"采集股票 {security_code} 数据失败: {str(e)}"
            logger.error(error_msg)
            if task_id:
                self._update_task_status(task_id, "failed", error_msg)
            return 0
    
    def _clean_stock_data(self, df):
        """清洗和标准化股票数据"""
        if df.empty:
            return df
        
        # 标准化列名
        df.columns = [col.lower() for col in df.columns]
        
        # 确保日期列格式正确
        if '日期' in df.columns:
            df.rename(columns={'日期': 'trade_date'}, inplace=True)
        
        # 标准化中文列名
        name_mapping = {
            '开盘': 'open',
            '收盘': 'close',
            '最高': 'high',
            '最低': 'low',
            '成交量': 'volume',
            '成交额': 'amount',
            '涨跌幅': 'change_pct',
            '换手率': 'turnover'
        }
        
        df.rename(columns=name_mapping, inplace=True)
        
        # 处理日期格式
        if 'trade_date' in df.columns:
            df['trade_date'] = pd.to_datetime(df['trade_date'])
        
        # 处理缺失值
        numeric_cols = ['open', 'close', 'high', 'low', 'volume', 'amount']
        for col in numeric_cols:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')
        
        # 删除交易量为0的无效数据
        if 'volume' in df.columns:
            df = df[df['volume'] > 0]
        
        # 保留需要的列
        keep_cols = [col for col in ['trade_date', 'open', 'high', 'low', 'close', 
                                     'volume', 'amount', 'turnover'] if col in df.columns]
        return df[keep_cols]
    
    def _save_daily_quotes(self, security_id, df):
        """保存日线数据到数据库"""
        if df.empty:
            return 0
        
        inserted_count = 0
        
        try:
            # 获取已存在的交易日期
            existing_dates = {
                date[0] for date in 
                self.db.query(DailyQuote.trade_date).filter(
                    DailyQuote.security_id == security_id
                ).all()
            }
            
            # 准备批量插入
            batch_records = []
            
            for _, row in df.iterrows():
                trade_date = row['trade_date'].date()
                
                # 跳过已存在的记录
                if trade_date in existing_dates:
                    continue
                
                daily_quote = DailyQuote(
                    security_id=security_id,
                    trade_date=trade_date,
                    open=row.get('open'),
                    high=row.get('high'),
                    low=row.get('low'),
                    close=row.get('close'),
                    adj_close=row.get('close'),  # 这里使用的是前复权数据
                    volume=row.get('volume'),
                    amount=row.get('amount'),
                    turnover=row.get('turnover')
                )
                
                batch_records.append(daily_quote)
                inserted_count += 1
                
                # 分批提交
                if len(batch_records) >= self.batch_size:
                    self.db.add_all(batch_records)
                    self.db.commit()
                    batch_records = []
            
            # 提交剩余记录
            if batch_records:
                self.db.add_all(batch_records)
                self.db.commit()
            
            logger.info(f"成功保存 {inserted_count} 条股票日线数据")
            return inserted_count
            
        except IntegrityError as e:
            self.db.rollback()
            logger.error(f"保存股票日线数据失败(完整性错误): {str(e)}")
            return 0
        except Exception as e:
            self.db.rollback()
            logger.error(f"保存股票日线数据失败: {str(e)}")
            raise
    def _update_task_status(self, task_id, status, error_message=None, result_count=None):
        """更新任务状态"""
        try:
            task = self.db.query(CollectionTask).filter(CollectionTask.task_id == task_id).first()
            if task:
                task.status = status
                if error_message:
                    task.error_message = error_message
                if result_count is not None:
                    task.result_count = result_count
                task.update_time = datetime.datetime.now()
                self.db.commit()
        except Exception as e:
            self.db.rollback()
            logger.error(f"更新任务状态失败: {str(e)}")

    def collect_chip_distribution(self, security_code, task_id=None):
        """
        采集股票筹码分布数据
        
        Args:
            security_code: 股票代码
            task_id: 任务ID
            
        Returns:
            成功保存的数据条数
        """
        # 更新任务状态
        if task_id:
            self._update_task_status(task_id, "running")
            
        # 获取证券信息
        security = self.db.query(SecurityInfo).filter(
            SecurityInfo.code == security_code,
            SecurityInfo.type == 'stock'
        ).first()
        
        if not security:
            error_msg = f"股票代码 {security_code} 不存在"
            logger.error(error_msg)
            if task_id:
                self._update_task_status(task_id, "failed", error_msg)
            return 0
            
        logger.info(f"开始采集股票 {security_code} 的筹码分布数据")
        
        try:
            # 调用AKShare接口获取筹码分布数据
            df = ak.stock_chip_distribution(symbol=security_code)
            
            if df.empty:
                logger.info(f"股票 {security_code} 没有可用的筹码分布数据")
                if task_id:
                    self._update_task_status(task_id, "success", None, 0)
                return 0
                
            # 清洗和转换数据
            df = self._clean_chip_data(df)
            
            # 保存到数据库
            inserted_count = self._save_chip_data(security.id, df)
            
            # 更新任务状态
            if task_id:
                self._update_task_status(task_id, "success", None, inserted_count)
                
            return inserted_count
            
        except Exception as e:
            error_msg = f"采集股票 {security_code} 筹码分布数据失败: {str(e)}"
            logger.error(error_msg)
            if task_id:
                self._update_task_status(task_id, "failed", error_msg)
            return 0
            
    def collect_capital_flow(self, security_code, task_id=None):
        """
        采集股票资金流数据
        
        Args:
            security_code: 股票代码
            task_id: 任务ID
            
        Returns:
            成功保存的数据条数
        """
        # 更新任务状态
        if task_id:
            self._update_task_status(task_id, "running")
            
        # 获取证券信息
        security = self.db.query(SecurityInfo).filter(
            SecurityInfo.code == security_code,
            SecurityInfo.type == 'stock'
        ).first()
        
        if not security:
            error_msg = f"股票代码 {security_code} 不存在"
            logger.error(error_msg)
            if task_id:
                self._update_task_status(task_id, "failed", error_msg)
            return 0
            
        logger.info(f"开始采集股票 {security_code} 的资金流数据")
        
        try:
            # 调用AKShare接口获取资金流数据
            df = ak.stock_capital_flow(symbol=security_code)
            
            if df.empty:
                logger.info(f"股票 {security_code} 没有可用的资金流数据")
                if task_id:
                    self._update_task_status(task_id, "success", None, 0)
                return 0
                
            # 清洗和转换数据
            df = self._clean_capital_flow_data(df)
            
            # 保存到数据库
            inserted_count = self._save_capital_flow_data(security.id, df)
            
            # 更新任务状态
            if task_id:
                self._update_task_status(task_id, "success", None, inserted_count)
                
            return inserted_count
            
        except Exception as e:
            error_msg = f"采集股票 {security_code} 资金流数据失败: {str(e)}"
            logger.error(error_msg)
            if task_id:
                self._update_task_status(task_id, "failed", error_msg)
            return 0
            
    def collect_financial_data(self, security_code, report_type="income", year=None, task_id=None):
        """
        采集股票财务数据
        
        Args:
            security_code: 股票代码
            report_type: 报表类型(income/balance/cashflow)
            year: 年份，默认为None表示最新年报
            task_id: 任务ID
            
        Returns:
            成功保存的数据条数
        """
        # 更新任务状态
        if task_id:
            self._update_task_status(task_id, "running")
        
        # 获取证券信息
        security = self.db.query(SecurityInfo).filter(
            SecurityInfo.code == security_code,
            SecurityInfo.type == 'stock'
        ).first()
        
        if not security:
            error_msg = f"股票代码 {security_code} 不存在"
            logger.error(error_msg)
            if task_id:
                self._update_task_status(task_id, "failed", error_msg)
            return 0
        
        logger.info(f"开始采集股票 {security_code} 的财务数据，报表类型: {report_type}")
        
        try:
            # 根据报表类型调用不同API
            if report_type == "income":
                df = ak.stock_financial_report_sina(stock=security_code, symbol="利润表")
            elif report_type == "balance":
                df = ak.stock_financial_report_sina(stock=security_code, symbol="资产负债表")
            elif report_type == "cashflow":
                df = ak.stock_financial_report_sina(stock=security_code, symbol="现金流量表")
            else:
                raise ValueError(f"不支持的报表类型: {report_type}")
            
            if df.empty:
                logger.info(f"股票 {security_code} 没有可用的财务数据")
                if task_id:
                    self._update_task_status(task_id, "success", None, 0)
                return 0
            
            # 清洗和转换数据
            df = self._clean_financial_data(df, report_type)
            
            # 筛选指定年份
            if year:
                df = df[df['report_date'].str.contains(str(year))]
            
            # 保存到数据库
            inserted_count = self._save_financial_data(security.id, report_type, df)
            
            # 更新任务状态
            if task_id:
                self._update_task_status(task_id, "success", None, inserted_count)
            
            return inserted_count
            
        except Exception as e:
            error_msg = f"采集股票 {security_code} 财务数据失败: {str(e)}"
            logger.error(error_msg)
            if task_id:
                self._update_task_status(task_id, "failed", error_msg)
            return 0
    
    def _clean_financial_data(self, df, report_type):
        """清洗和标准化财务数据"""
        if df.empty:
            return df
        
        # 标准化列名
        df.columns = [col.strip() for col in df.columns]
        
        # 添加报表类型
        df['report_type'] = report_type
        
        # 处理日期格式
        if '报告日期' in df.columns:
            df.rename(columns={'报告日期': 'report_date'}, inplace=True)
            df['report_date'] = pd.to_datetime(df['report_date'])
        
        # 处理数值类型
        numeric_cols = [col for col in df.columns if col not in ['report_date', 'report_type']]
        for col in numeric_cols:
            df[col] = pd.to_numeric(df[col], errors='coerce')
        
        return df
    
    def _save_financial_data(self, security_id, report_type, df):
        """保存财务数据到数据库"""
        if df.empty:
            return 0
        
        inserted_count = 0
        
        try:
            # 获取已存在的记录
            existing_records = {
                (record.report_date, record.item_name) for record in
                self.db.query(FinancialData.report_date, FinancialData.item_name).filter(
                    FinancialData.security_id == security_id,
                    FinancialData.report_type == report_type
                ).all()
            }
            
            # 准备批量插入
            batch_records = []
            
            for _, row in df.iterrows():
                report_date = row['report_date'].date()
                
                # 遍历所有财务项目
                for col in df.columns:
                    if col in ['report_date', 'report_type']:
                        continue
                    
                    item_name = col
                    item_value = row[col]
                    
                    # 跳过已存在的记录
                    if (report_date, item_name) in existing_records:
                        continue
                    
                    financial_data = FinancialData(
                        security_id=security_id,
                        report_type=report_type,
                        report_date=report_date,
                        item_name=item_name,
                        item_value=item_value,
                        update_time=datetime.datetime.now()
                    )
                    
                    batch_records.append(financial_data)
                    inserted_count += 1
                    
                    # 分批提交
                    if len(batch_records) >= self.batch_size:
                        self.db.add_all(batch_records)
                        self.db.commit()
                        batch_records = []
            
            # 提交剩余记录
            if batch_records:
                self.db.add_all(batch_records)
                self.db.commit()
            
            logger.info(f"成功保存 {inserted_count} 条财务数据")
            return inserted_count
            
        except IntegrityError as e:
            self.db.rollback()
            logger.error(f"保存财务数据失败(完整性错误): {str(e)}")
            return 0
        except Exception as e:
            self.db.rollback()
            logger.error(f"保存财务数据失败: {str(e)}")
            raise



# 测试函数
def test_stock_collector():
    """测试股票数据采集器功能"""
    collector = StockCollector()
    
    # 1. 测试获取股票列表
    try:
        stock_list = collector.get_all_stock_list()
        print(f"成功获取 {len(stock_list)} 只股票基本信息")
        print(stock_list.head())
    except Exception as e:
        print(f"获取股票列表失败: {str(e)}")
    
    # 2. 测试采集股票基本信息
    try:
        new_count, update_count = collector.collect_stock_info()
        print(f"新增股票: {new_count}, 更新股票: {update_count}")
    except Exception as e:
        print(f"采集股票基本信息失败: {str(e)}")
    
    # 3. 测试采集日线数据
    if stock_list is not None and not stock_list.empty:
        # 选择第一只股票进行测试
        test_code = stock_list.iloc[0]['code']
        try:
            # 采集最近30天数据
            end_date = datetime.datetime.now().strftime('%Y%m%d')
            start_date = (datetime.datetime.now() - datetime.timedelta(days=30)).strftime('%Y%m%d')
            
            count = collector.collect_daily_quote(test_code, start_date, end_date)
            print(f"股票 {test_code} 成功采集 {count} 条日线数据")
        except Exception as e:
            print(f"采集股票 {test_code} 日线数据失败: {str(e)}")
        
        # 4. 测试采集财务数据
        try:
            # 测试利润表
            income_count = collector.collect_financial_data(test_code, "income")
            print(f"股票 {test_code} 成功采集 {income_count} 条利润表数据")
            
            # 测试资产负债表
            balance_count = collector.collect_financial_data(test_code, "balance")
            print(f"股票 {test_code} 成功采集 {balance_count} 条资产负债表数据")
            
            # 测试现金流量表
            cashflow_count = collector.collect_financial_data(test_code, "cashflow")
            print(f"股票 {test_code} 成功采集 {cashflow_count} 条现金流量表数据")
        except Exception as e:
            print(f"采集股票 {test_code} 财务数据失败: {str(e)}")
        
        # 5. 测试采集筹码分布数据
        try:
            chip_count = collector.collect_chip_distribution(test_code)
            print(f"股票 {test_code} 成功采集 {chip_count} 条筹码分布数据")
        except Exception as e:
            print(f"采集股票 {test_code} 筹码分布数据失败: {str(e)}")
        
        # 6. 测试采集资金流数据
        try:
            flow_count = collector.collect_capital_flow(test_code)
            print(f"股票 {test_code} 成功采集 {flow_count} 条资金流数据")
        except Exception as e:
            print(f"采集股票 {test_code} 资金流数据失败: {str(e)}")


if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 执行测试
    test_stock_collector()